import os
import json
import pandas as pd
from typing import Dict, List, Tuple, Optional
from pathlib import Path
import logging


class StockAdjustment:
  """股票复权处理类"""

  def __init__(self, raw_data_dir: str, xr_data_dir: str, ipo_data_dir: str, output_dir: str):
    """
        初始化复权处理器

        Args:
            raw_data_dir: 原始数据目录路径
            xr_data_dir: 除权数据目录路径
            ipo_data_dir: IPO数据目录路径
            output_dir: 输出目录路径
        """
    self.raw_data_dir = Path(raw_data_dir)
    self.xr_data_dir = Path(xr_data_dir)
    self.ipo_data_dir = Path(ipo_data_dir)
    self.output_dir = Path(output_dir)

    # 创建输出目录
    self.output_dir.mkdir(parents=True, exist_ok=True)
    self.forward_dir = self.output_dir / "forward"
    self.backward_dir = self.output_dir / "backward"
    self.forward_dir.mkdir(exist_ok=True)
    self.backward_dir.mkdir(exist_ok=True)

    # 设置日志
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
    self.logger = logging.getLogger(__name__)

  def load_stock_data(self, file_path: Path) -> List[Dict]:
    """加载股票数据文件"""
    try:
      with open(file_path, 'r', encoding='utf-8') as f:
        return json.load(f)
    except Exception as e:
      self.logger.error(f"加载文件 {file_path} 失败: {e}")
      return []

  def get_stock_codes(self) -> List[str]:
    """获取所有股票代码"""
    raw_files = list(self.raw_data_dir.glob("*.json"))
    return [f.stem for f in raw_files]

  def calculate_adjustment_factors(self, xr_data: List[Dict], price_data: List[Dict]) -> Tuple[pd.Series, pd.Series]:
    """
        计算前复权和后复权因子

        Args:
            xr_data: 除权数据
            price_data: 价格数据

        Returns:
            forward_factors: 前复权因子
            backward_factors: 后复权因子
        """
    # 创建价格数据DataFrame并按日期排序
    price_df = pd.DataFrame(price_data)
    price_df['date'] = pd.to_datetime(price_df['date'], format='%Y%m%d')
    price_df = price_df.sort_values('date').set_index('date')

    # 创建除权数据DataFrame并按日期排序
    xr_df = pd.DataFrame(xr_data)
    xr_df['date'] = pd.to_datetime(xr_df['date'], format='%Y%m%d')
    xr_df = xr_df.sort_values('date')

    # 初始化复权因子
    forward_factors = pd.Series(1.0, index=price_df.index)
    backward_factors = pd.Series(1.0, index=price_df.index)

    # 计算每股分红和转增
    xr_df['per_share_dividend'] = xr_df['dividend'] / 10.0  # 每股分红
    xr_df['per_share_convert'] = xr_df['convert'] / 10.0  # 每股转增

    current_forward_factor = 1.0
    current_backward_factor = 1.0

    # 按日期顺序处理除权事件
    for _, xr_event in xr_df.iterrows():
      xr_date = xr_event['date']

      # 找到除权日前一天的收盘价
      prev_dates = price_df[price_df.index < xr_date]
      if len(prev_dates) == 0:
        continue

      prev_date = prev_dates.index.max()
      prev_close = price_df.loc[prev_date, 'close']

      # 计算调整比例
      # 前复权: (前收盘价 - 每股分红) / (前收盘价 * (1 + 每股转增))
      adjust_ratio = (prev_close - xr_event['per_share_dividend']) / (prev_close * (1 + xr_event['per_share_convert']))

      # 更新前复权因子
      mask_forward = price_df.index <= xr_date
      forward_factors[mask_forward] = forward_factors[mask_forward] * adjust_ratio

      # 更新后复权因子
      mask_backward = price_df.index >= xr_date
      backward_factors[mask_backward] = backward_factors[mask_backward] / adjust_ratio

    return forward_factors, backward_factors

  def adjust_price_data(self, price_data: List[Dict], factors: pd.Series, adjustment_type: str) -> List[Dict]:
    """
        应用复权因子调整价格数据

        Args:
            price_data: 原始价格数据
            factors: 复权因子
            adjustment_type: 复权类型 ('forward' 或 'backward')

        Returns:
            复权后的价格数据
        """
    # 创建DataFrame
    df = pd.DataFrame(price_data)
    df['date'] = pd.to_datetime(df['date'], format='%Y%m%d')
    df = df.sort_values('date')

    # 确保日期索引匹配
    df = df.set_index('date')
    aligned_factors = factors.reindex(df.index, method='ffill').fillna(1.0)

    # 应用复权因子
    price_columns = ['open', 'high', 'low', 'close', 'last_close']
    for col in price_columns:
      if col in df.columns:
        df[f'{col}_{adjustment_type}'] = df[col] * aligned_factors

    # 只保留复权列，移除原始列
    adjusted_columns = [f'{col}_{adjustment_type}' for col in price_columns]
    result_columns = adjusted_columns + ['volume', 'turnover_rate']
    
    # 移除不在result_columns中的列
    for col in df.columns:
        if col not in result_columns and col not in adjusted_columns:
            df = df.drop(columns=[col])
    
    # 重命名复权列，去掉后缀
    rename_map = {f'{col}_{adjustment_type}': col for col in price_columns}
    df = df.rename(columns=rename_map)

    # 重置索引并转换日期格式
    result_df = df.reset_index()
    result_df['date'] = result_df['date'].dt.strftime('%Y%m%d')

    return result_df.to_dict('records')

  def process_single_stock(self, stock_code: str, output_types: List[str] = ['forward', 'backward']) -> bool:
    """
        处理单只股票的复权数据

        Args:
            stock_code: 股票代码
            output_types: 输出类型列表 ['forward', 'backward'] 或其中之一

        Returns:
            处理是否成功
        """
    try:
      # 加载数据
      raw_file = self.raw_data_dir / f"{stock_code}.json"
      xr_file = self.xr_data_dir / f"{stock_code}.json"
      ipo_file = self.ipo_data_dir / f"{stock_code}.json"

      if not raw_file.exists():
        self.logger.warning(f"原始数据文件不存在: {raw_file}")
        return False

      price_data = self.load_stock_data(raw_file)
      if not price_data:
        return False

      xr_data = []
      if xr_file.exists():
        xr_data = self.load_stock_data(xr_file)

      ipo_data = []
      if ipo_file.exists():
        ipo_data = self.load_stock_data(ipo_file)

      # 计算复权因子
      forward_factors, backward_factors = self.calculate_adjustment_factors(xr_data, price_data)

      # 生成输出数据
      results = {}

      if 'forward' in output_types:
        forward_data = self.adjust_price_data(price_data, forward_factors, 'forward')
        results['forward'] = forward_data

      if 'backward' in output_types:
        backward_data = self.adjust_price_data(price_data, backward_factors, 'backward')
        results['backward'] = backward_data

      # 保存结果到不同文件夹
      for adj_type, data in results.items():
        if adj_type == 'forward':
            output_file = self.forward_dir / f"{stock_code}.json"
        elif adj_type == 'backward':
            output_file = self.backward_dir / f"{stock_code}.json"
        else:
            output_file = self.output_dir / f"{stock_code}_{adj_type}.json"
            
        with open(output_file, 'w', encoding='utf-8') as f:
          json.dump(data, f, ensure_ascii=False, indent=2)

      self.logger.info(f"成功处理股票 {stock_code}: {output_types}")
      return True

    except Exception as e:
      self.logger.error(f"处理股票 {stock_code} 失败: {e}")
      return False

  def process_all_stocks(self, output_types: List[str] = ['forward', 'backward'], max_workers: int = None):
    """
        处理所有股票的复权数据

        Args:
            output_types: 输出类型列表
            max_workers: 最大工作线程数（None表示使用单线程）
        """
    stock_codes = self.get_stock_codes()
    self.logger.info(f"开始处理 {len(stock_codes)} 只股票")

    success_count = 0
    failed_stocks = []

    if max_workers and max_workers > 1:
      # 使用多线程处理
      from concurrent.futures import ThreadPoolExecutor, as_completed

      with ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_to_stock = {
          executor.submit(self.process_single_stock, code, output_types): code
          for code in stock_codes
        }

        for future in as_completed(future_to_stock):
          stock_code = future_to_stock[future]
          try:
            if future.result():
              success_count += 1
            else:
              failed_stocks.append(stock_code)
          except Exception as e:
            self.logger.error(f"股票 {stock_code} 处理异常: {e}")
            failed_stocks.append(stock_code)
    else:
      # 单线程处理
      for stock_code in stock_codes:
        if self.process_single_stock(stock_code, output_types):
          success_count += 1
        else:
          failed_stocks.append(stock_code)

    self.logger.info(f"处理完成: 成功 {success_count}, 失败 {len(failed_stocks)}")
    if failed_stocks:
      self.logger.warning(f"失败的股票: {failed_stocks}")


def generate_adjusted_data(
    raw_data_dir: str,
    xr_data_dir: str,
    ipo_data_dir: str,
    output_dir: str,
    adjustment_types: List[str] = ['forward', 'backward'],
    max_workers: int = None
):
  """
    生成复权数据的主函数

    Args:
        raw_data_dir: 原始数据目录路径
        xr_data_dir: 除权数据目录路径
        ipo_data_dir: IPO数据目录路径
        output_dir: 输出目录路径
        adjustment_types: 复权类型列表 ['forward', 'backward'] 或其中之一
        max_workers: 最大工作线程数
    """
  # 验证输入参数
  valid_types = ['forward', 'backward']
  for adj_type in adjustment_types:
    if adj_type not in valid_types:
      raise ValueError(f"不支持的复权类型: {adj_type}，支持的类型: {valid_types}")

  processor = StockAdjustment(raw_data_dir, xr_data_dir, ipo_data_dir, output_dir)
  processor.process_all_stocks(adjustment_types, max_workers)


# 使用示例
if __name__ == "__main__":
  # 示例用法
  generate_adjusted_data(
    raw_data_dir="./行情",
    xr_data_dir="./除息",
    ipo_data_dir="./IPO",
    output_dir="./复权",
    adjustment_types=['forward', 'backward'],  # 同时生成前复权和后复权
    max_workers=4  # 使用4个线程并行处理
  )

  # 或者只生成前复权
  # generate_adjusted_data(
  #     raw_data_dir="./raw_data",
  #     xr_data_dir="./xr_data",
  #     ipo_data_dir="./ipo_data",
  #     output_dir="./adjusted_data",
  #     adjustment_types=['forward']  # 只生成前复权
  # )